import os
import cv2
import numpy as np

# 根目录
image_root = '/home/lhx/pcbsyn/data_zoo/pcb_truecolor_subcpjrsbg/images'
depth_root = '/home/lhx/pcbsyn/data_zoo/pcb_truecolor_subcpjrsbg/depths'
rgbd_root = '/home/lhx/pcbsyn/data_zoo/pcb_truecolor_subcpjrsbg/rgbd'

# 遍历 train, val, test 三个子文件夹
for subset in ['train', 'val', 'test']:
    image_dir = os.path.join(image_root, subset)
    depth_dir = os.path.join(depth_root, subset)
    rgbd_dir = os.path.join(rgbd_root, subset)

    # 创建输出文件夹
    os.makedirs(rgbd_dir, exist_ok=True)

    # 遍历图像文件
    for root, _, files in os.walk(image_dir):
        for file in files:
            if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
                # 构建对应路径
                rel_path = os.path.relpath(os.path.join(root, file), image_dir)
                rgb_path = os.path.join(image_dir, rel_path)
                depth_path = os.path.join(depth_dir, rel_path)
                output_path = os.path.join(rgbd_dir, rel_path)

                # 确保输出目录存在
                os.makedirs(os.path.dirname(output_path), exist_ok=True)

                # 加载RGB和Depth图像
                rgb = cv2.imread(rgb_path, cv2.IMREAD_COLOR)         # (H, W, 3)
                depth = cv2.imread(depth_path, cv2.IMREAD_GRAYSCALE) # (H, W)

                if rgb is None or depth is None:
                    print(f"跳过无法读取的图像：{rgb_path} 或 {depth_path}")
                    continue

                # 确保深度图大小与RGB一致
                if rgb.shape[:2] != depth.shape[:2]:
                    print(f"大小不一致：{rgb_path} vs {depth_path}")
                    continue

                # 合并为4通道
                depth_expanded = np.expand_dims(depth, axis=2)       # (H, W, 1)
                rgbd = np.concatenate((rgb, depth_expanded), axis=2) # (H, W, 4)

                # 保存为png
                cv2.imwrite(output_path, rgbd)

print("融合完成！RGBD图像已保存到 rgbd 文件夹。")
